#!/usr/bin/env python
# encoding: utf-8
"""
rowfile.py

Created by Colin Raffel, Jieun Oh, and Isaac Wang for MUS 422 Final Project
"""

from pcmfile import *
import numpy as np
import scipy.io.wavfile as wav
np.set_printoptions(precision=3,suppress=True)
import quantize
import sys
import noisecode
import mp3wav
import os
import huffmanCode
import codec
import eyeD3
import psychoac
import trainData
import pickle
import warnings
warnings.simplefilter('ignore')

def Encode(inFileName, bitRate, useTonal=False):
    """
    Encode a wav file as a row.mp3
    args:
        inFileName - wav file
        bitRate - bitRate of row.mp3
    """
    # Find difference (error) between wav and mp3
    (diffName, mp3File) = mp3wav.findDifference(inFileName, bitRate)
    inFile = PCMFile(diffName)

    # Setup coding params
    codingParams = inFile.OpenForReading()
    setupCodingParams(codingParams)

    # How many samples read in per block?
    blockSize = 1024
    # Store this value in the codingParams object
    codingParams.nSamplesPerBlock = blockSize
    # Get sampling frequency
    Fs = codingParams.sampleRate

    # Create one noiseTonalSeparator and noiseDecoder class per channel
    noiseTonalSeparators = []
    
    for iCh in range(codingParams.nChannels):
        # noiseTonalSeparator args: fftSize, Fs, noiseLeak, tonalLeak, tonalSmear, tonalPower
        noiseTonalSeparators.append( noisecode.noiseTonalSeparator( blockSize, Fs, .95, .7, 1.0, 10.0 ) )

    # Arrays that will be stored in ID3 tag
    noiseArr = [] # Huffman Code first
    mantissaArr = [] # Huffman Code first
    scaleArr = []
    bitAllocArr = []
    overallScaleFactorArr = []
    
    print "Encoding error/noise using PAC files..."
    
    # Create output files
    while True:
        # Read in data
        data=inFile.ReadDataBlock(codingParams)
        # If we are at the end of the file, break out
        if not data: break
    
        fullBlockData=[]
        # For each channel...
        for iCh in range(codingParams.nChannels):
            # Tick in the data for each channel
            noiseTonalSeparators[iCh].tick(data[iCh])
            # Get noise levels (here we are doing quantization)
            noiseLevels = noiseTonalSeparators[iCh].noiseCode()
            noiseArr = np.append(noiseArr,noiseLevels)
            
            if useTonal:
                fullBlockData.append(noiseTonalSeparators[iCh].getTonal())
            else:
                fullBlockData.append(np.concatenate((codingParams.priorBlock[iCh], data[iCh])))
        
        # Current pass's data is next pass's prior block data
        codingParams.priorBlock = data
        
        # Encode the full block of multi-channel data using PAC files
        (scaleFactor, bitAlloc, mantissa, overallScaleFactor) = codec.Encode(fullBlockData,codingParams)
        mantissaArr.append(mantissa)
        scaleArr.append(scaleFactor)
        bitAllocArr.append(bitAlloc)
        overallScaleFactorArr.append(overallScaleFactor)

    # Huffman Coding
    print "Huffman coding noise envelope..."
    (Code2Symbol, hcMantissaVecs) = trainData.trainNoise(noiseArr)
    # huffmanMants = trainData.trainData(mantissaArr)

    # ID3 Tagging
    print "Inserting extra data into ID3 TXXX tag..."
    tag = eyeD3.Tag()
    tag.link(mp3File)
    tag.setVersion(eyeD3.ID3_V2)
    
    tag.addUserTextFrame('huffmanNoise',pickle.dumps((Code2Symbol, hcMantissaVecs))) # Dump huffman as a string to be stored in ID3 tag
    tag.addUserTextFrame('huffmanMants',pickle.dumps(mantissaArr))
    tag.addUserTextFrame('scaleArr',pickle.dumps(scaleArr))
    tag.addUserTextFrame('bitAllocArr',pickle.dumps(bitAllocArr))
    tag.addUserTextFrame('overallScaleFactorArr',pickle.dumps(overallScaleFactorArr))
    tag.addUserTextFrame('codingParams',pickle.dumps(codingParams))
    
    tag.update()
    
    # os.remove(diffName)
    return mp3File

def Decode(mp3File):
    """docstring for Decode"""
    
    # Grab data from ID3 tags
    print "Decoding error/noise from ID3 TXXX tag..."
    tag = eyeD3.Tag()
    tag.link(mp3File)
    for frame in tag.getUserTextFrames():
        if frame.description == 'huffmanNoise':
            (Code2Symbol, hcMantissaVecs) = pickle.loads(str(frame.text))
        elif frame.description == 'huffmanMants':
            huffmanMants = pickle.loads(str(frame.text))
        elif frame.description == 'scaleArr':
            scaleArr = pickle.loads(str(frame.text))
        elif frame.description == 'bitAllocArr':
            bitAllocArr = pickle.loads(str(frame.text))
        elif frame.description == 'overallScaleFactorArr':
            overallScaleFactorArr = pickle.loads(str(frame.text))
        elif frame.description == 'codingParams':
            codingParams = pickle.loads(str(frame.text))
    
    # Prepare coding params
    codingParams.bitsPerSample = 16
    overlapAndAdd = []
    for iCh in range(codingParams.nChannels): overlapAndAdd.append( np.zeros(codingParams.nMDCTLines, dtype=np.float64) )
    codingParams.overlapAndAdd=overlapAndAdd
    
    # Prepare noise parameters
    noiseDecoders = []
    noiseLevels = [] 
    
    for iCh in range(codingParams.nChannels):
        # noiseDecoder args: fftSize, Fs, noiseSmear
        noiseDecoders.append( noisecode.noiseDecoder( codingParams.nSamplesPerBlock, codingParams.sampleRate, 1.0 ) )
    
    # Un-Huffmanize
    # print Code2Symbol
    # print hcMantissaVecs
    for i in xrange(len(hcMantissaVecs)):
        for j in xrange(len(hcMantissaVecs[0])):
            # print 'i:',i,'j:',j
            # print hcMantissaVecs[i,j]
            noiseLevels.append(Code2Symbol[str(hcMantissaVecs[i,j])])
    
    # return noiseLevels
    
    # Prepare outFile
    outFileName = mp3File[:-4] + '.error' + '.wav'
    outFile= PCMFile(outFileName)
    outFile.OpenForWriting(codingParams)
    
    # print 'noiseLevels:',noiseLevels,len(noiseLevels)
    
    # Recreate error
    for i in xrange(len(scaleArr)):
        data = []
        for iCh in xrange(codingParams.nChannels):
            data.append(np.array([],dtype=np.float64))  # add location for this channel's data
            # Decode the unpacked data for this channel, overlap-and-add first half, and append it to the data array (saving other half for next overlap-and-add)
            decodedData = codec.Decode(list(scaleArr[i][0]),list(bitAllocArr[i][0]),huffmanMants[i][0], overallScaleFactorArr[i][0],codingParams)
            data[iCh] = 0.009 * np.concatenate( (data[iCh],np.add(codingParams.overlapAndAdd[iCh],decodedData[:codingParams.nMDCTLines]) ) )  # data[iCh] is overlap-and-added data
            
            # Add generated noise to data
            data[iCh] = data[iCh] + 0.005 * noiseDecoders[iCh].noiseDecode(noiseLevels[(2*i + iCh) * 25:(2*i + iCh + 1) * 25])
            codingParams.overlapAndAdd[iCh] = decodedData[codingParams.nMDCTLines:]  # save other half for next pass
        outFile.WriteDataBlock(data,codingParams)
    if codingParams.overlapAndAdd:
        outFile.WriteDataBlock(codingParams.overlapAndAdd,codingParams)
    
    outFile.Close(codingParams)
    
    # Add error + noise to original mp3 as wav
    mp3wav.synthesizeWav(mp3File,outFileName)
    os.remove(outFileName)
    

def setupCodingParams(codingParams):
    """Setup codingParams before reading data"""
    codingParams.nMDCTLines = 1024
    codingParams.nScaleBits = 3
    codingParams.nMantSizeBits = 2
    codingParams.targetBitsPerSample = 0.1
    # tell the PCM file how large the block size is
    codingParams.nSamplesPerBlock = codingParams.nMDCTLines
    
    # make sure that the number of samples in the file is a multiple of the number of MDCT half-blocksize, otherwise zero pad as needed
    if not codingParams.numSamples%codingParams.nMDCTLines:
        codingParams.numSamples += (codingParams.nMDCTLines
                    - codingParams.numSamples%codingParams.nMDCTLines) # zero padding for partial final PCM block
    # also add in the delay block for the second pass w/ the last half-block
    codingParams.numSamples+= codingParams.nMDCTLines  # due to the delay in processing the first samples on both sides of the MDCT block
    # create a ScaleFactorBand object to be used by the encoding process and write its info to header
    sfBands=psychoac.ScaleFactorBands(psychoac.AssignMDCTLinesFromFreqLimits(codingParams.nMDCTLines,codingParams.sampleRate))
    codingParams.sfBands=sfBands

    # start w/o all zeroes as prior block of unencoded data for other half of MDCT block
    priorBlock = []
    overlapAndAdd = []
    for iCh in range(codingParams.nChannels):
        priorBlock.append(np.zeros(codingParams.nMDCTLines,dtype=np.float64) )
        overlapAndAdd.append( np.zeros(codingParams.nMDCTLines, dtype=np.float64) )
    codingParams.priorBlock = priorBlock
    codingParams.overlapAndAdd = overlapAndAdd
    return

if __name__ == "__main__":
    print "#################################################################################"
    print "############################ ROW.mp3 Encoder/Decoder ############################"
    print "#################################################################################"
    print " "
    # Get input arguments (if there are any)
    inFileName = ''
    outFileName = ''
    bitRate = 320
    useTonal = False

    if len(sys.argv) > 2:
        # Get filenames
        inFileName = sys.argv[1]
        bitRate = sys.argv[2]
        # useTonal = sys.argv[3]

    # If the inpt args are invalid, exit
    else:
        while True:
            inFileName = raw_input('Input file name (must end in .wav): ')
            if inFileName != '' and os.path.isfile(inFileName):
                break
            else:
                print 'Invalid input file name'
        bR = raw_input('BitRate [default is 320]: ')
        if bR == "":
            bitRate = 320
        else:
            bitRate = int(bR)
    
    print "Encoding", inFileName, "using a bit rate of", bitRate
    mp3File = Encode(inFileName, bitRate, useTonal)
    print "Decoding", mp3File
    Decode(mp3File)
    